#include "detection_vision/tensorrt/impl/utils.hpp"
#include "detection_vision/impl/bbox2d.hpp"
#include <iostream>
#include <cstdio>
#include <cstdlib>

namespace detection
{ 

__device__ float Logist(float data){ return 1./(1. + exp(-data)); }


__global__ void CenterNetDecoder( float *hm,  float *reg,  float *wh, 
    float *output, const int w, const int h,  const int classes, const int kernel_size,
     const float score_threshold, const int down_ratio)
{
    int C = threadIdx.z;
    int B = blockIdx.z ;
    int col = blockIdx.x * blockDim.x  + threadIdx.x;
    int row = blockIdx.y * blockDim.y  + threadIdx.y;
    if(col >= w || row >= h) return;
    // printf("%d",resCount);
    char *data_head = (char *)output + B *( w * h * 5 * sizeof(float) + sizeof(unsigned int));
       
    int idx  = row * w + col + C * w * h;
    // int reg_idx = row * w + col;
    float* p_data_reg = reg + B * 2 * w * h; 
    float* p_data_wh = wh + B * 2 * w * h;
    float* hm_cur_batch = hm + B * w * h * classes;
    float hm_logist = Logist(hm_cur_batch[idx]);
    int feature_map_size =  w * h;
    int offset = -(kernel_size - 1) / 2;
    int  l, m;
    if(hm_logist >= (score_threshold))
    {
        float max = -1;
        int max_index = 0;
        for (l = 0; l < kernel_size; ++l)
            for (m = 0; m < kernel_size; ++m) {
                int cur_x = offset + l + col;
                int cur_y = offset + m + row;
                int cur_index = cur_y * w + cur_x + feature_map_size * C;
                int valid = (cur_x >= 0 && cur_x < w && cur_y >= 0 && cur_y < h);
                float val = (valid != 0) ? Logist(hm_cur_batch[cur_index]) : -1;
                max_index = (val > max) ? cur_index : max_index;
                max = (val > max) ? val : max;
            }
        if(idx == max_index)
        {
            unsigned int* count_p = (unsigned int*)data_head;
            int resCount = (int) atomicAdd(count_p, 1);
            char *data = ((char *) data_head )+ sizeof(unsigned int) + resCount * sizeof(float) * 6;
            int  X_Y_x = row * w + col;
            int X_y_y = X_Y_x + feature_map_size;
            int x = col + 2*p_data_reg[X_Y_x];
            int y = row + 1.5*p_data_reg[X_y_y];
            int _w = p_data_wh[X_Y_x];
            int _h = p_data_wh[X_y_y];

            float* bbox_ptr = (float*)(data);
            *bbox_ptr = C;
            bbox_ptr++;
            *bbox_ptr = (x )* down_ratio - (_w * down_ratio)/2;
            bbox_ptr++;
            *bbox_ptr = (y ) * down_ratio - (_h * down_ratio)/2;
            bbox_ptr++;
            *bbox_ptr = _w * down_ratio;
            bbox_ptr++;
            *bbox_ptr = _h * down_ratio;
            bbox_ptr++;
            *bbox_ptr = hm_logist;
        }

    }
}

void CenterNetDetDecoderGPU( float *hm,  float *reg,  float *wh, const int B, 
                            const int cnn_input_width, const int cnn_input_height, const int classes,
                            const int kernel_size,const int down_ratio, const float score_threshold, 
                            float *output)
{
    int block_size = 8;
    unsigned int threads_x = cnn_input_width/down_ratio > 0 ? block_size : cnn_input_width/down_ratio;
    unsigned int threads_y = cnn_input_height/down_ratio > 0 ? block_size : cnn_input_height/down_ratio;
    unsigned int blocks_x = cnn_input_width/down_ratio/threads_x > 0 ? cnn_input_width/down_ratio/threads_x : 1; 
    unsigned int blocks_y = cnn_input_height/down_ratio/threads_y > 0 ? cnn_input_height/down_ratio/threads_y : 1; 

    dim3 grids = {blocks_x, blocks_y, (unsigned int)B};
    dim3 threads = {threads_x, threads_y, (unsigned int)classes};
    CenterNetDecoder<<<grids, threads>>>(hm, reg, wh, output, cnn_input_width/down_ratio,
                                        cnn_input_height / down_ratio, classes, kernel_size,
                                        score_threshold, down_ratio);
}

}